{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy import symbols, cos, sin, pi, simplify, trigsimp, expand_trig, pprint\n",
    "from sympy.matrices import Matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "q1, q2, q3, q4, q5, q6= symbols('q1:7')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pi\n"
     ]
    }
   ],
   "source": [
    "print(pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def pose(theta, alpha, a, d):\n",
    "\n",
    "  r11, r12 = cos(theta), -sin(theta)\n",
    "  r23, r33 = -sin(alpha), cos(alpha)\n",
    "  r21 = sin(theta) * cos(alpha)\n",
    "  r22 = cos(theta) * cos(alpha)\n",
    "  r31 = sin(theta) * sin(alpha)\n",
    "  r32 = cos(theta) * sin(alpha)\n",
    "  y = -d * sin(alpha)\n",
    "  z = d * cos(alpha)\n",
    "    \n",
    "  T = Matrix([\n",
    "    [r11, r12, 0.0, a],\n",
    "    [r21, r22, r23, y],\n",
    "    [r31, r32, r33, z],\n",
    "    [0.0, 0.0, 0.0, 1]\n",
    "  ])\n",
    "  \n",
    "  T = simplify(T)\n",
    "\n",
    "  return T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_pretty_matrix(T):\n",
    "\n",
    "  T11, T12, T13, T14 = str(T[0, 0]), str(T[0, 1]), str(T[0, 2]), str(T[0, 3])\n",
    "  T21, T22, T23, T24 = str(T[1, 0]), str(T[1, 1]), str(T[1, 2]), str(T[1, 3])\n",
    "  T31, T32, T33, T34 = str(T[2, 0]), str(T[2, 1]), str(T[2, 2]), str(T[2, 3])\n",
    "  T41, T42, T43, T44 = str(T[3, 0]), str(T[3, 1]), str(T[3, 2]), str(T[3, 3])\n",
    "\n",
    "  s1 = '|{:20.20}|{:20.20}|{:20.20}|{:30.30}|'.format(T11, T12, T13, T14)\n",
    "  s2 = '|{:20.20}|{:20.20}|{:20.20}|{:30.30}|'.format(T21, T22, T23, T24)\n",
    "  s3 = '|{:20.20}|{:20.20}|{:20.20}|{:30.30}|'.format(T31, T32, T33, T34)\n",
    "  s4 = '|{:20.20}|{:20.20}|{:20.20}|{:30.30}|'.format(T41, T42, T43, T44)\n",
    "    \n",
    "  print('\\n', s1, '\\n', s2, '\\n', s3, '\\n', s4, '\\n')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "T01\n",
      "\n",
      " |cos(q1)             |-sin(q1)            |0.0                 |0                             | \n",
      " |sin(q1)             |cos(q1)             |0                   |0                             | \n",
      " |0                   |0                   |1                   |0.750000000000000             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T12\n",
      "\n",
      " |sin(q2)             |cos(q2)             |0.0                 |0.350000000000000             | \n",
      " |0                   |0                   |1                   |0                             | \n",
      " |cos(q2)             |-sin(q2)            |0                   |0                             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T23\n",
      "\n",
      " |cos(q3)             |-sin(q3)            |0.0                 |1.25000000000000              | \n",
      " |sin(q3)             |cos(q3)             |0                   |0                             | \n",
      " |0                   |0                   |1                   |0                             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T34\n",
      "\n",
      " |cos(q4)             |-sin(q4)            |0.0                 |-0.0540000000000000           | \n",
      " |0                   |0                   |1                   |1.50000000000000              | \n",
      " |-sin(q4)            |-cos(q4)            |0                   |0                             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T45\n",
      "\n",
      " |cos(q5)             |-sin(q5)            |0.0                 |0                             | \n",
      " |0                   |0                   |-1                  |0                             | \n",
      " |sin(q5)             |cos(q5)             |0                   |0                             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T56\n",
      "\n",
      " |cos(q6)             |-sin(q6)            |0.0                 |0                             | \n",
      " |0                   |0                   |1                   |0                             | \n",
      " |-sin(q6)            |-cos(q6)            |0                   |0                             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n",
      "T6g\n",
      "\n",
      " |1                   |0                   |0.0                 |0                             | \n",
      " |0                   |1                   |0                   |0                             | \n",
      " |0                   |0                   |1                   |0.303000000000000             | \n",
      " |0.0                 |0.0                 |0.0                 |1                             | \n",
      "\n"
     ]
    }
   ],
   "source": [
    "d90 = pi / 2 \n",
    "\n",
    "T01 = pose(q1, 0, 0, 0.75)\n",
    "T12 = pose(q2 - d90, -d90, 0.35, 0)\n",
    "T23 = pose(q3, 0, 1.25, 0)\n",
    "T34 = pose(q4, -d90, -0.054, 1.5)\n",
    "T45 = pose(q5, d90, 0, 0)\n",
    "T56 = pose(q6, -d90, 0, 0)\n",
    "T6g = pose(0, 0, 0, 0.303)\n",
    "\n",
    "print(\"T01\")\n",
    "print_pretty_matrix(T01)\n",
    "\n",
    "print(\"T12\")\n",
    "print_pretty_matrix(T12)\n",
    "\n",
    "print(\"T23\")\n",
    "print_pretty_matrix(T23)\n",
    "\n",
    "print(\"T34\")\n",
    "print_pretty_matrix(T34)\n",
    "\n",
    "print(\"T45\")\n",
    "print_pretty_matrix(T45)\n",
    "\n",
    "print(\"T56\")\n",
    "print_pretty_matrix(T56)\n",
    "\n",
    "print(\"T6g\")\n",
    "print_pretty_matrix(T6g)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sin(a)*sin(b) + cos(a)*cos(b)\n"
     ]
    }
   ],
   "source": [
    "a, b = symbols('a b')\n",
    "print(expand_trig(cos(a - b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cos(a - b)\n"
     ]
    }
   ],
   "source": [
    "print(trigsimp(sin(a)*sin(b) + cos(a)*cos(b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "T03 = simplify(T01 * T12 * T23) \n",
    "R03 = T03[:3, :3]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "T03\n",
      "\n",
      " |sin(q2 + q3)*cos(q1)|cos(q1)*cos(q2 + q3)|-sin(q1)            |(1.25*sin(q2) + 0.35)*cos(q1) | \n",
      " |sin(q1)*sin(q2 + q3)|sin(q1)*cos(q2 + q3)|cos(q1)             |(1.25*sin(q2) + 0.35)*sin(q1) | \n",
      " |cos(q2 + q3)        |-sin(q2 + q3)       |0                   |1.25*cos(q2) + 0.75           | \n",
      " |0                   |0                   |0                   |1                             | \n",
      "\n",
      "\n",
      "R03\n",
      "⎡sin(q₂ + q₃)⋅cos(q₁)  cos(q₁)⋅cos(q₂ + q₃)  -sin(q₁)⎤\n",
      "⎢                                                    ⎥\n",
      "⎢sin(q₁)⋅sin(q₂ + q₃)  sin(q₁)⋅cos(q₂ + q₃)  cos(q₁) ⎥\n",
      "⎢                                                    ⎥\n",
      "⎣    cos(q₂ + q₃)         -sin(q₂ + q₃)         0    ⎦\n"
     ]
    }
   ],
   "source": [
    "print()\n",
    "print(\"T03\")\n",
    "print_pretty_matrix(T03)\n",
    "\n",
    "print()\n",
    "print(\"R03\")\n",
    "pprint(R03)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "R03\n",
      "Matrix([\n",
      "[sin(q2 + q3)*cos(q1), cos(q1)*cos(q2 + q3), -sin(q1)],\n",
      "[sin(q1)*sin(q2 + q3), sin(q1)*cos(q2 + q3),  cos(q1)],\n",
      "[        cos(q2 + q3),        -sin(q2 + q3),        0]])\n"
     ]
    }
   ],
   "source": [
    "print(\"R03\")\n",
    "print(R03)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡-sin(q₄)⋅sin(q₆) + cos(q₄)⋅cos(q₅)⋅cos(q₆)  -sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄\n",
      "⎢                                                                             \n",
      "⎢             sin(q₅)⋅cos(q₆)                             -sin(q₅)⋅sin(q₆)    \n",
      "⎢                                                                             \n",
      "⎢-sin(q₄)⋅cos(q₅)⋅cos(q₆) - sin(q₆)⋅cos(q₄)  sin(q₄)⋅sin(q₆)⋅cos(q₅) - cos(q₄)\n",
      "⎢                                                                             \n",
      "⎣                    0                                           0            \n",
      "\n",
      ")⋅cos(q₅)  -sin(q₅)⋅cos(q₄)  -0.054⎤\n",
      "                                   ⎥\n",
      "               cos(q₅)        1.5  ⎥\n",
      "                                   ⎥\n",
      "⋅cos(q₆)   sin(q₄)⋅sin(q₅)     0   ⎥\n",
      "                                   ⎥\n",
      "                  0            1   ⎦\n"
     ]
    }
   ],
   "source": [
    "T36 = T34*T45*T56\n",
    "pprint(T36)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⎡-sin(q₄)⋅sin(q₆) + cos(q₄)⋅cos(q₅)⋅cos(q₆)  -sin(q₄)⋅cos(q₆) - sin(q₆)⋅cos(q₄\n",
      "⎢                                                                             \n",
      "⎢             sin(q₅)⋅cos(q₆)                             -sin(q₅)⋅sin(q₆)    \n",
      "⎢                                                                             \n",
      "⎣-sin(q₄)⋅cos(q₅)⋅cos(q₆) - sin(q₆)⋅cos(q₄)  sin(q₄)⋅sin(q₆)⋅cos(q₅) - cos(q₄)\n",
      "\n",
      ")⋅cos(q₅)  -sin(q₅)⋅cos(q₄)⎤\n",
      "                           ⎥\n",
      "               cos(q₅)     ⎥\n",
      "                           ⎥\n",
      "⋅cos(q₆)   sin(q₄)⋅sin(q₅) ⎦\n"
     ]
    }
   ],
   "source": [
    "R36 = T36[:3, :3]\n",
    "pprint(R36)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [],
   "source": [
    "R36f = Matrix([[-sin(q4)*sin(q6) + cos(q4)*cos(q5)*cos(q6), -sin(q4)*cos(q6) - sin(q6)*cos(q4)*cos(q5), -sin(q5)*cos(q4)],\n",
    "  [                           sin(q5)*cos(q6),                           -sin(q5)*sin(q6),          cos(q5)],\n",
    "  [-sin(q4)*cos(q5)*cos(q6) - sin(q6)*cos(q4),  sin(q4)*sin(q6)*cos(q5) - cos(q4)*cos(q6),  sin(q4)*sin(q5)]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "True\n"
     ]
    }
   ],
   "source": [
    "print(R36 == R36f)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  },
  "widgets": {
   "state": {},
   "version": "1.1.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
